from __future__ import print_function
from __future__ import division

import sys

from sklearn2pmml import PMMLPipeline, sklearn2pmml, make_pmml_pipeline
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.feature_extraction.text import TfidfTransformer
from sklearn.linear_model import LogisticRegression
import pickle
from sklearn_pandas import DataFrameMapper
from sklearn2pmml.feature_extraction.text import Splitter


if __name__ == '__main__':
    cv_file = sys.argv[1]
    tfidf_file = sys.argv[2]
    lg_file = sys.argv[3]

    cv = pickle.load(open(cv_file, 'rb'))
    tfidf = pickle.load(open(tfidf_file, 'rb'))
    lg = pickle.load(open(lg_file, 'rb'))

    CountVectorizer().tokenizer = Splitter()

    # steps = [CountVectorizer(), TfidfTransformer(), LogisticRegression()]
    steps = [('Count_vectorizer', cv), ('tfidf', tfidf), ('classifier', lg)]
    pipeline = PMMLPipeline(steps)

    pipeline_pmml = make_pmml_pipeline(pipeline)

    sklearn2pmml(pipeline_pmml, '/Users/hardy/lr.pmml')

